from datetime import datetime
from config.settings import get_settings
from get_logger import get_data_logger, get_server_logger, get_error_logger
from check_video import CheckVideo, CheckResults
from extends.redis import redis_connection

import cv2
import os
import numpy as np
import time
import redis
import orjson

data_logger = get_data_logger()
error_logger = get_error_logger()
server_logger = get_server_logger()

x1, y1 = 0, 120
x2, y2 = 1910, 441

x = min(x1, x2)
y = min(y1, y2)
width = abs(x2 - x1)
height = abs(y2 - y1)

# 定义颜色的阈值范围（可以根据实际情况调整）
green_lower = np.array([40, 40, 40])
green_upper = np.array([80, 255, 255])

red_lower1 = np.array([0, 50, 50])
red_upper1 = np.array([5, 255, 255])
red_lower2 = np.array([175, 50, 50])
red_upper2 = np.array([180, 255, 255])

blue_lower = np.array([100, 50, 50])
blue_upper = np.array([130, 255, 255])

tracks = [
    "901",
    "902",
    "903",
    "904",
    "905",
    "911",
    "912"
]

settings = get_settings()



class SmartWatchDog:
    check_video = CheckVideo()
    redis_conn: redis.Redis = None
    video: cv2.VideoCapture = None 
    
    results = {
        "Capture_delay": "",
        "Timestamp": "",
        "Remark": "",
        "Valid": True,
        "Loco": {"status": {}},
    }
    start_time: float = time.time()
    found_detection_updates = False
    previous_areas = [] # 上一次识别的矩形区域大小列表
    exit = False


    def update_results(self, status, check_res:CheckResults):
        """更新响应数据"""
        self.results["Timestamp"] = datetime.now().strftime(
        "%m/%d/%Y %H:%M:%S.%f")

        self.results["Loco"]["status"] = status
        self.results["Capture_delay"] = check_res.get("datetime_ocr", "")

        # remark为空时,才填写心跳
        if settings.ENABLED_HEART_BEAT:
            if time.time() - self.start_time > settings.HEART_INTERVAL:
                self.results.setdefault("Remark", "Keep Alive - Normal")
                self.found_detection_updates = True
        return self.results
        
    def start(self):
        if settings.MODE == "test":
            test_video_path = os.path.join(
                settings.BASE_DIR, "video", "output_20240607092354.avi")
            self.video = cv2.VideoCapture(test_video_path)
        else:
            self.video = cv2.VideoCapture(0)
            self.video.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
            self.video.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
            print(f"VideoIsOpen:{self.video.isOpened()}")

        return_value, frame = self.video.read()

        if not return_value:
            # Error
            print("cannot grab frame.")
            return None
        else:
            # 跳过5帧
            for _ in range(5):
                _, _ = self.video.read()

        previous_results = {}
        min_area = 2500
        max_area = 4000
        while return_value:
            return_value, frame = self.video.read()
            
            frame1 = frame[y:y+height, x:x+width]
            check_res = self.check_video.thread_start_checkers(frame)

            ## init data
            self.found_detection_updates = False
            self.results["Remark"] = ""

            # 查找矩形
            rectangles = []
            areas = []  # 用于存储每个矩形区域的面积

            # 动态识别大小阈值
            gray = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
            blurred = cv2.GaussianBlur(gray, (5, 5), 0)
            edges = cv2.Canny(blurred, 50, 150)
            kernel = np.ones((5, 5), np.uint8)
            edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
            contours, _ = cv2.findContours(
                edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            for contour in contours:
                perimeter = cv2.arcLength(contour, True)
                approx = cv2.approxPolyDP(contour, 0.02 * perimeter, True)
                if len(approx) == 4:
                    x3, y3, w3, h3 = cv2.boundingRect(contour)
                    area = cv2.contourArea(contour)
                    if area > min_area and area < max_area:
                        rectangles.append((x3, y3, w3, h3))
                        areas.append(area)

            # 保留大小相似的
            if len(rectangles) and (areas != self.previous_areas):
                self.previous_areas = areas
                filtered_areas_index = self.filter_similar_elements(
                    areas, percentage=0.05)
                if filtered_areas_index and len(filtered_areas_index) != len(rectangles):
                    rectangles = [rectangles[i] for i in filtered_areas_index]
                    filter_area = [areas[i] for i in filtered_areas_index]
                    average_area = self.average(filter_area)
                    # 减少识别大小
                    if (average_area - 200) >= 0:
                        min_area = average_area - 200

                    if (average_area + 200) <= 5000:
                        max_area = average_area + 200

            elif len(rectangles) < 7:
                # 扩大识别大小
                if min_area >= 500:
                    min_area -= 500

                if max_area <= 5000:
                    max_area += 500
                    continue

            # 排序
            rectangles_sorted = sorted(rectangles, key=lambda rect: rect[0])

            status = {track:"" for track in tracks}

            status = self.detect_and_update_status_from_colors(frame1=frame1, rectangles_sorted=rectangles_sorted, status=status)
            results = self.update_results(status=status, check_res=check_res)
            if status != previous_results or self.found_detection_updates:
                previous_results = status
                self.start_time = time.time()
                data_logger.info(results)
                if self.redis_conn:
                    self.redis_conn.set(settings.DATA_RESULT_KEY, orjson.dumps(results))

            if len(rectangles_sorted) == 7:
                cv2.imshow(
                    "Hung Hom Stabling Sidings Smart Watch Dog",
                    cv2.resize(
                        frame1,
                        (800, 120)
                    )
                )
            else:
                cv2.imshow(
                    "Hung Hom Stabling Sidings Smart Watch Dog",
                    cv2.resize(
                        frame,
                        (800, 410)
                    )
                )
            key = cv2.waitKey(20)

            if key == 113:
                break
            
            if key == 32:
                # press Space to capture image (Space ASCII value: 32)
                now = datetime.now()
                dt_string = now.strftime("%Y%m%d%H%M%S")
                img_name = \
                    "opencv_frame_{}.png".format(dt_string)

                cv2.imwrite(img_name, frame)

    def detect_success(self, rectangles_sorted, frame1, status:dict) -> dict:
        """画面检测成功画框和修改路段数据"""
        self.results["Valid"] = True
        for index, (x3, y3, w3, h3) in enumerate(rectangles_sorted):
            cv2.rectangle(frame1, (x3, y3),
                        (x3 + w3, y3 + h3), (0, 255, 10), 5)
            left_top_x = x3 + int(w3 / 8)
            left_top_y = y3 + int(h3 / 6)
            right_bottom_x = x3 + w3 - int(w3 / 8)
            right_bottom_y = y3 + h3 - int(h3 / 6)

            # # 提取这些点的颜色
            color1 = frame1[left_top_y, left_top_x]  # 左上角点
            color2 = frame1[right_bottom_y, right_bottom_x]  # 右下角点

            # 将颜色从BGR转换到HSV空间
            hsv_color1 = cv2.cvtColor(
                np.uint8([[color1]]), cv2.COLOR_BGR2HSV)[0][0]
            hsv_color2 = cv2.cvtColor(
                np.uint8([[color2]]), cv2.COLOR_BGR2HSV)[0][0]

            if blue_lower[0] <= hsv_color1[0] <= blue_upper[0] and \
                    blue_lower[1] <= hsv_color1[1] <= blue_upper[1] and \
                    blue_lower[2] <= hsv_color1[2] <= blue_upper[2] and \
                    blue_lower[0] <= hsv_color2[0] <= blue_upper[0] and \
                    blue_lower[1] <= hsv_color2[1] <= blue_upper[1] and \
                    blue_lower[2] <= hsv_color2[2] <= blue_upper[2]:
                status[tracks[index]] = "R"
            elif green_lower[0] <= hsv_color1[0] <= green_upper[0] and \
                    green_lower[1] <= hsv_color1[1] <= green_upper[1] and \
                    green_lower[2] <= hsv_color1[2] <= green_upper[2] and \
                    green_lower[0] <= hsv_color2[0] <= green_upper[0] and \
                    green_lower[1] <= hsv_color2[1] <= green_upper[1] and \
                    green_lower[2] <= hsv_color2[2] <= green_upper[2]:
                status[tracks[index]] = "N"
            elif (red_lower1[0] <= hsv_color1[0] <= red_upper1[0] or red_lower2[0] <= hsv_color1[0] <= red_upper2[0]) and \
                red_lower1[1] <= hsv_color1[1] <= red_upper1[1] and \
                red_lower1[2] <= hsv_color1[2] <= red_upper1[2] and \
                (red_lower1[0] <= hsv_color2[0] <= red_upper1[0] or red_lower2[0] <= hsv_color2[0] <= red_upper2[0]) and \
                red_lower1[1] <= hsv_color2[1] <= red_upper1[1] and \
                    red_lower1[2] <= hsv_color2[2] <= red_upper1[2]:
                status[tracks[index]] = "OC"
        return status

    def detect_fail(self, rectangles_sorted, frame1):
        """画面异常, 只给他画框, 并修改提示"""
        self.results["Valid"] = False
        self.results["Remark"] = "Display Abnormality"
        for (x3, y3, w3, h3) in rectangles_sorted:
            cv2.rectangle(frame1, (x3, y3),(x3 + w3, y3 + h3), (0, 255, 10), 5)

    def detect_and_update_status_from_colors(self, frame1, rectangles_sorted, status:dict):
        """识别方块的颜色, 并更新状态"""
        if len(rectangles_sorted) == 7:
            status = self.detect_success(rectangles_sorted, frame1, status)
        else:
            self.detect_fail(rectangles_sorted, frame1)
        return status

    def filter_similar_elements(self, elements, percentage=0.05):
        """保留大小相似的元素"""
        server_logger.info(f"Filter Similar Elements Input: {elements}")
        thresholds = [element * percentage for element in elements]
        filtered_elements_index = []

        for i, elem in enumerate(elements):
            elem_threshold = thresholds[i]
            similar_elem_found = False

            for value in elements:
                if abs(elem - value) <= elem_threshold and elem != value:
                    similar_elem_found = True
                    break

            if similar_elem_found:
                filtered_elements_index.append(int(i))

        return filtered_elements_index

    def average(self, lst):
        """求平均值"""
        server_logger.info(f"Average Input: {lst}")
        if not lst:
            return 3300
        average = sum(lst) / len(lst)
        return int(average)

    def video_close(self):
        if self.video is not None:
            self.video.release()

    def close(self):
        self.check_video.close()
        redis_connection.close()
        server_logger.info("redis close")

        self.video_close()
        server_logger.info("Video Close")
        
        cv2.destroyAllWindows()
        server_logger.info("Exit")

        self.exit = True
    def __enter__(self):
        server_logger.info("redis connection ...")
        while not self.redis_conn:
            try:
                self.redis_conn = redis_connection.get_redis_connection()
            except Exception as e:
                server_logger.error("redis connection error: %s", str(e))
                time.sleep(3)
        server_logger.info("redis connection success")
        server_logger.info("SmartWatchDog Start")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if not self.exit:
            self.close()

if __name__ == "__main__":
    with SmartWatchDog() as smart_watch_dog:
        try:
            smart_watch_dog.start()
        except (KeyboardInterrupt) as e:
            error_logger.error(e)
        finally:
            smart_watch_dog.close()
            
